Anime Face Stable Diffusion Model

What are stable diffusion models

Stable diffusion models are state-of-art deep generative models in the machine learning field. They are capable of generating photo-realistic images from latent space and have been gaining significant attention in the past several years. Many famous generative AI websites DALL-E 2, Midjourney, etc are based on the technique of stable diffusion models.

What is the goal of this notebook

In this notebook, I tried to built an *unconditional* denoising diffusion model from scratch and used the model to generate anime faces.

Anime faces generated from latent space:

My codes and results are shown in below.

References:
Denoising Diffusion Probabilistic Models
Improved Denoising Diffusion Probabilistic Models
How stable diffusion work

Github:
Sung-Soo Kim's implementation
Niels Rogge & Kashif Rasul's implementation

Datasets for Training

I used Anime_Faces_Dataset_128 scraped from kaggle which consists of around 100K high quality anime faces images (128*128 - white background - same pose)

import torch
import torchvision
import matplotlib.pyplot as plt
import numpy as np
from torch import nn
from torch.optim import Adam, SGD, AdamW
from torchvision import transforms
from torch.utils.data import DataLoader
from torchinfo import summary
from PIL import Image
from tqdm import tqdm

""" Plots some samples from the dataset """
nsamples = 20
ncols = 5
data = torchvision.datasets.ImageFolder(root=r'/Users/mwchang/Downloads/faces')
indices = torch.randint(low=0, high=len(data), size=(nsamples,))
samples = torch.utils.data.Subset(data, indices)
plt.figure(figsize=(15,15))
for i, img in enumerate(samples):
    plt.subplot(int(nsamples/ncols) + 1, ncols, i + 1)
    plt.imshow(img[0])
del(data)

Forward and Reverse Diffusion Process

In a denoise diffusion model, the noise-corrupted input data is treated as the initial state. The model then applies a series of denoising steps, where each step reduces the noise level and refines the data.

During the denoising process, the model estimates the posterior distribution of the clean data given the current state and noise level. It leverages the inherent patterns and structure in the data to guide the denoising process. By iteratively applying these denoising steps, the model gradually removes the noise and converges towards a cleaner representation of the original data.

Ref websites:
Denoising Diffusion Probabilistic Models
Improved Denoising Diffusion Probabilistic Models
How stable diffusion work
Positional Encoding

Implementation of Diffusion

class Diffusion:
    """ This class implements the diffusion process described in the papers:
        https://arxiv.org/abs/2006.11239
        https://arxiv.org/abs/2102.09672?ref=assemblyai.com
    """

    def __init__(
            self,
            noise_steps: int = 1000,
            beta_1: float = 1e-4,
            beta_t: float = 0.02,
            img_size: int = 64,
            device: str = 'cpu') -> None:

        self.noise_steps = noise_steps  # T
        self.img_size = img_size
        self.device = device

        # Gaussian noise to the data according to a variance schedule β1, . . . , βT:
        self.beta = torch.linspace(beta_1, beta_t, noise_steps)  # shape: (noise_steps,)

        # alpha = 1 − βt
        self.alpha = 1.0 - self.beta  # shape: (noise_steps,)

        # alpha_bar = ∏t′=1,...,t (1 − αt′)
        self.alpha_bar = torch.cumprod(self.alpha, dim=0)  # shape: (noise_steps,)

        # Move to the target device:
        self.beta = self.beta.to(self.device)   # shape: (noise_steps,)
        self.alpha = self.alpha.to(self.device)  # shape: (noise_steps,)
        self.alpha_bar = self.alpha_bar.to(self.device)  # shape: (noise_steps,)

    def sample_timesteps(self, n: int = 1) -> torch.Tensor:
        """ Sample n timesteps uniformly from [1, noise_steps]"""
        return torch.randint(low=1, high=self.noise_steps, size=(n,)).long().to(self.device)  # shape: (n,)

    def noise_image_at_time_t(self, x0: torch.Tensor, t: int or torch.Tensor) -> tuple[torch.Tensor]:
        """ Forward diffusion process:
        Add Gaussian noise to the image according to:

            xt(x0, epsilon) = √ᾱt x0 + √(1 − ᾱt) epsilon, where epsilon ~ N(0, I)

        Parameters:
            x0: image_to_be_noised, shape: (N, C, H, W)
            t: time step, shape: (N, )
        """
        if type(t) == int:
            # Convert t to a tensor:
            t = (torch.ones(x0.shape[0]) * t).long().to(self.device)   # shape: (N, )

        # Get the alpha_bar at time step t:
        alpha_bar_t = self.alpha_bar[t][:, None, None, None]  # shape: (N, 1, 1, 1)

        # Generate Gaussian noise:
        epsilon = torch.randn_like(x0)  # shape: (N, C, H, W); normal distribution with mean=0, std=1

        # Add Gaussian noise to the data:
        xt = torch.sqrt(alpha_bar_t) * x0 + torch.sqrt(1.0 - alpha_bar_t) * epsilon  # shape: (N, C, H, W)
        return xt, epsilon

    def denoise_at_time_t(self, model: nn.Module, xt: torch.Tensor, t: int) -> torch.Tensor:
        """ Denoise at time step t

         xt-1 = 1/√αt * (xt - (1 - αt)/√(1 - ᾱt) * e(xt, t)) + √βt * z,

        where:
            xt-1 is the denoised image at time step t-1,
            xt is the noised image at time step t,
            z ~ N(0, I), and
            e(xt, t) is the predicted noise at time step t, x

        Parameters:
            model: the model used to predict the noise
            xt: the noised image at time step t, shape: (N, C, H, W)
            t: time step, shape: (N, )

        Returns:
            xt_1: the denoised image at time step t-1, shape: (N, C, H, W)
        """

        if t > 1:
            z = torch.randn(xt.shape).to(self.device)  # shape: (N, C, H, W)
        else:
            z = torch.zeros(xt.shape).to(self.device)  # shape: (N, C, H, W)

        # Convert t to a tensor:
        t = (torch.ones(xt.shape[0]) * t).long().to(self.device)   # shape: (N, )

        # Get the predicted noise:
        predicted_noise = model(xt, t)  # shape: (N, C, H, W)

        # Get the alpha at time step t:
        alpha_t = self.alpha[t][:, None, None, None]  # shape: (N, 1, 1, 1)

        # Get the alpha_bar at time step t:
        alpha_bar_t = self.alpha_bar[t][:, None, None, None]  # shape: (N, 1, 1, 1)

        # Get the beta at time step t:
        beta_t = self.beta[t][:, None, None, None]  # shape: (N, 1, 1, 1)

        # Denoise:
        xt_1 = 1 / torch.sqrt(alpha_t) * (
                xt - ((1 - alpha_t) / (torch.sqrt(1 - alpha_bar_t))) * predicted_noise
        ) + torch.sqrt(beta_t) * z  # shape: (N, C, H, W)

        return xt_1  # x at time step t-1

    @torch.no_grad()
    def sample(self, model: nn.Module, n: int) -> torch.Tensor:
        """ Sample n images from the model.
         This is the implementation of Algorithm 2 in the paper: https://arxiv.org/abs/2006.11239

        Parameters:
            model: the model used to predict the noise
            n: number of images to be sampled
         """
        # Set the model to evaluation mode:
        model.eval()

        # Sample noise: xT ~ N(0, I)
        x = torch.randn((n, 3, self.img_size, self.img_size)).to(self.device)  # shape: (N, C, H, W)

        # Denoise from time step T to 1:
        for t in tqdm(reversed(range(1, self.noise_steps)), position=0):   # t = T, T-1, ..., 1
            x = self.denoise_at_time_t(model=model, xt=x, t=t)  # x at time step t-1

        # Clamp the values to [-1, 1]
        x0 = x  # x at time step 0

        # # Set the model back to training mode:
        model.train()
        return x0  # shape: (N, C, H, W)

Positional Encoding

Diffusion models typically involve modeling the noise distribution at each step of the diffusion process. The noise distribution may vary over time, and one needs to incorporate time-dependent information or time-related features into the model the model to capture temporal variations in the noise characteristics. Such time series data can be considered as sequential information. It can be described by the position of an entity in a sequence so that each position is assigned a unique representation.

Positional encoding is a technique first proposed in a transformer paper: Attention Is All You Need to address the lack of sequential information in transformer models. The technique is applied for encoding time sequential information as well.

Ref websites:
https://kazemnejad.com/blog/transformer_architecture_positional_encoding/
https://machinelearningmastery.com/a-gentle-introduction-to-positional-encoding-in-transformer-models-part-1/

Implementation of PositionalEncoder

class PositionalEncoder(nn.Module):
    """ This class implements the Position Encoder from the paper:
        https://arxiv.org/abs/1706.03762

        Ref website:
        https://kazemnejad.com/blog/transformer_architecture_positional_encoding/
        https://machinelearningmastery.com/a-gentle-introduction-to-positional-encoding-in-transformer-models-part-1/
    """

    _pem = None  # positional embedding matrix

    def __init__(self, seq_len: int, emb_dim: int) -> None:
        """
        Parameters:
        out_features: the number of output features of the layer
        seq_len: the length of the sequence or total time steps,
        emb_dim: the dimension of the embedding vector
        """

        super().__init__()
        self.seq_len = seq_len
        self.emb_dim = emb_dim

        # To save memory, we will only create the positional embedding matrix once
        if self._pem is None or self._pem.shape != (seq_len, emb_dim):
            self._pem = self.create_sinusoidal_pos_emb(seq_len, emb_dim)

    @staticmethod
    def create_sinusoidal_pos_emb(seq_len: int, dim: int) -> torch.Tensor:
        """
        PE(pos, 2i) = sin(pos / 10000^(2i/dim))
        PE(pos, 2i+1) = cos(pos / 10000^(2i/dim))

        where:
            pos is the time position,
            i is the index of the dimension,
            dim is the dimension of the embedding vector.

        """
        pem = torch.zeros(seq_len, dim)  # shape: (seq_length, dim)
        pos = torch.arange(0, seq_len).unsqueeze(1)  # shape: (seq_length, 1)
        i = torch.arange(0, dim, 2)  # shape: (dim / 2, )
        pem[:, 0::2] = torch.sin(pos / 10000 ** (2 * i / dim))
        pem[:, 1::2] = torch.cos(pos / 10000 ** (2 * i / dim))
        return pem

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        """ Forward pass of the Positional Embedding Layer
         Parameters:
            t: the time step, shape: (N, ), where N is the batch size,

         Returns:
             time embedding vectors, shape: (N, dim)
         """
        # Move the positional embedding matrix to the same device as t:
        if self._pem.device != t.device:
            self._pem = self._pem.to(t.device)

        # Get the positional embedding vector for each time step:
        return self._pem[t, :]  # shape: (N, dim)

UNet architecture and diffusion models

UNet has been widely adopted in various computer vision tasks. It was introduced by Olaf Ronneberger, Philipp Fischer, and Thomas Brox in their 2015 paper titled U-Net: Convolutional Networks for Biomedical Image Segmentation. The name "UNet" comes from the U-shaped architecture of the network, which consists of a contracting path (downsampling path) and an expanding path (upsampling path).

In this Notebook, the UNet architecture is adopted as the denoiser or generator component within the diffusion model framework. At each step, the UNet takes the noisy image as input and produces a denoised version. This denoised version is then used as the input for the next step, progressively reducing the noise level. The UNet can leverage its ability to capture both local and global features to enhance the denoising process. The network looks like:

Ref websites:
What is UNET?
UNet — Line by Line Explanation.

Implementation of Unet

class Block(nn.Module):
    """ This class implements the Double Convolutional Layers """

    def __init__(self, in_channels: int, out_channels: int, emb_dim: int = 256, residual=True) -> None:
        super().__init__()
        self.residual = residual

        # First Convolutional Layer
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)  # shape: (N, out_channels, H, W)
        self.batchnorm1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()

        # Time Embedding Layer
        self.time_emb = nn.Sequential(
            nn.Linear(emb_dim, out_channels),
            nn.ReLU(),
        )

        # Second Convolutional Layer
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)  # shape: (N, out_channels, H, W)
        self.batchnorm2 = nn.BatchNorm2d(out_channels)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """  shape of x: (N, in_channels, H, W),
             shape of t: (N, emb_dim)
        """
        # Pass the input tensor through the first convolutional layer:
        x = self.conv1(x)  # shape: (N, in_channels, H, W) -> (N, out_channels, H, W)
        x = self.relu(x)  # shape: (N, out_channels, H, W) -> (N, out_channels, H, W)
        x = self.batchnorm1(x)  # shape: (N, out_channels, H, W) -> (N, out_channels, H, W)

        # Add the time embedding vector:
        t = self.time_emb(t)  # shape: (N, emb_dim) -> (N, out_channels)
        t = t.unsqueeze(-1).unsqueeze(-1)  # shape: (N, out_channels, 1, 1)

        # Add the time embedding vector to the input tensor:
        x = x + t  # shape: (N, out_channels, H, W)  -> (N, out_channels, H, W)

        # Pass the input tensor through the second convolutional layer:
        x = self.conv2(x)  # shape: (N, out_channels, H, W) -> (N, out_channels, H, W)

        # Add the residual:
        if self.residual:
            x = x + self.relu(x)  # shape:  (N, out_channels, H, W) -> (N, out_channels, H, W)
        else:
            x = self.relu(x)  # shape:  (N, out_channels, H, W) -> (N, out_channels, H, W)

        x = self.batchnorm2(x)  # shape:  (N, out_channels, H, W) -> (N, out_channels, H, W)

        return x  # shape: (N, out_channels, H, W)


class Down(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, emb_dim: int = 256, residual: bool = True):
        super().__init__()

        # self.down = nn.MaxPool2d(2)  # shape: (N, in_channels, H/2, W/2)
        self.down = nn.Conv2d(in_channels, in_channels, kernel_size=4, stride=2, padding=1)
        self.conv = Block(in_channels, out_channels, emb_dim, residual)  # shape: (N, out_channels, H/2, W/2)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """  shape of x: (N, in_channels, H, W),
             shape of t: (N, emb_dim)"""
        x = self.down(x)  # shape: (N, in_channels, H, W) -> (N, in_channels, H/2, W/2)
        x = self.conv(x, t)  # shape: (N, in_channels, H/2, W/2) -> (N, out_channels, H/2, W/2)
        return x # shape: (N, out_channels, H/2, W/2) + (N, out_channels, 1, 1) -> (N, out_channels, H/2, W/2)


class Up(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, emb_dim: int = 256, residual: bool = True):
        super().__init__()

        # self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)  # shape: (N, in_channels, H*2, W*2)
        self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)  # shape: (N, in_channels, H*2, W*2)
        self.conv = Block(2 * out_channels, out_channels, emb_dim, residual)  # shape: (N, out_channels, H*2, W*2)

    def forward(self, x: torch.Tensor, skip_x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """  shape of x: (N, in_channels, H, W),
             shape of skip_x: (N, in_channels, H*2, W*2)
             shape of t: (N, emb_dim)
        """
        x = self.up(x)  # shape: (N, in_channels, H, W) -> (N, out_channels, H*2, W*2)
        x = torch.cat([x, skip_x], dim=1)  # shape: (N, in_channels + in_channels, H*2, W*2) -> (N, 2*in_channels, H*2, W*2)
        x = self.conv(x, t)  # shape: (N, 2*in_channels, H*2, W*2) -> (N, out_channels, H*2, W*2)
        return x  # shape: (N, out_channels, H*2, W*2) -> (N, out_channels, H*2, W*2)


class UNet(nn.Module):
    def __init__(self,
                 in_channels: int = 3,
                 out_channels: int = 3,
                 seq_len: int = 1000,
                 emb_dim: int = 256):
        super().__init__()

        self.time_emb = nn.Sequential(
            PositionalEncoder(seq_len, emb_dim),  # shape: (N, dim)
            nn.Linear(emb_dim, emb_dim),  # shape: (N, dim)
            nn.ReLU(),  # shape: (N, dim)
        )  # shape: (N, dim)

        self.inc = Block(in_channels, 64, emb_dim=emb_dim)  # (N, 64, img_size, img_size)
        self.down1 = Down(64, 128, emb_dim=emb_dim)  # (N, 128, img_size/2, img_size/2)
        self.down2 = Down(128, 256, emb_dim=emb_dim)  # (N, 256, img_size/4, img_size/4)
        self.down3 = Down(256, 512, emb_dim=emb_dim)  # (N, 512, img_size/8, img_size/8)
        self.down4 = Down(512, 512, emb_dim=emb_dim)  # (N, 512, img_size/16, img_size/16)
        self.bot1 = Block(512, 1024, emb_dim=emb_dim)  # (N, 1024, img_size/16, img_size/16)
        self.up1 = Up(1024, 512, emb_dim=emb_dim)  # (N, 512, img_size/8, img_size/8)
        self.up2 = Up(512, 256, emb_dim=emb_dim)  # (N, 256, img_size/4, img_size/4))
        self.up3 = Up(256, 128, emb_dim=emb_dim)  # (N, 128, img_size/2, img_size/2)
        self.up4 = Up(128, 64, emb_dim=emb_dim)  # (N, 64, img_size, img_size)
        self.outc = Block(64, out_channels, emb_dim=emb_dim)  # (N, out_channels, img_size, img_size)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """ shape of x (N, in_channels, H, W),
            shape of t (N, )
         """
        # Initial:
        t = self.time_emb(t)  # shape: (N, emb_dim)
        xi1 = self.inc(x, t)  # shape: (N, in_channels, H, W) -> (N, 64, H, W)

        # Down:
        xd1 = self.down1(xi1, t)  # shape: (N, 64, H, W) -> (N, 128, H/2, W/2)
        xd2 = self.down2(xd1, t)  # shape: (N, 128, H/2, W/2) -> (N, 256, H/4, W/4)
        xd3 = self.down3(xd2, t)  # shape: (N, 256, H/4, W/4) -> (N, 512, H/8, W/8)
        xd4 = self.down4(xd3, t)  # shape: (N, 512, H/8, W/8) -> (N, 512, H/16, W/16)

        # Bottom:
        xb1 = self.bot1(xd4, t)  # shape: (N, 512, H/16, W/16) -> (N, 1024, H/16, W/16)

        # Up:
        x = self.up1(xb1, skip_x=xd3, t=t)  # shape: (N, 1024, H/16, W/16) -> (N, 512, H/8, W/8)
        x = self.up2(x, skip_x=xd2, t=t)  # shape: (N, 512, H/8, W/8) -> (N, 256, H/4, W/4)
        x = self.up3(x, skip_x=xd1, t=t)  # shape: (N, 256, H/4, W/4) -> (N, 128, H/2, W/2)
        x = self.up4(x, skip_x=xi1, t=t)  # shape: (N, 128, H/2, W/2) -> (N, 64, H, W)

        # Output
        x = self.outc(x, t)  # shape: (N, 64, H, W) -> (N, out_channels, H, W)
        return x

Helper functions

def plot_tensor_images(images: torch.Tensor,  path: str=None, **kwargs):
    images = images.clamp(-1, 1)  # clamp the values to [-1, 1]
    images = (images + 1) / 2.0  # mean = 0, std = 1 -> mean = 0.5, std = 0.5
    images = (images * 255).type(torch.uint8)  # convert x0 to [0, 255] and type: uint8

    # Create a grid of images:
    grid = torchvision.utils.make_grid(images, **kwargs)  # shape: (C, H, W)
    grid = grid.permute(1, 2, 0)  # shape: (H, W, C)
    
    # Convert to PIL image:
    img = Image.fromarray(grid.detach().cpu().numpy())
    
    plt.imshow(img)
    
    if path:
        img.save(path)

def save_checkpoint(model, optimizer, epoch, path=None):
    if path is None:
        path = f'model_epoch{epoch}.chk'
    
    torch.save({
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'epoch': epoch
    }, path)
    
    
def load_checkpoint(model, optimizer, path):
    checkpoint = torch.load(path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']
    return model, optimizer, epoch

Hyper parameters

T = 1000
IMG_SIZE = 64
IN_CHANNELS = 3
BATCH_SIZE = 128
TIME_EBD_DIM = 256
SUBSET = 20000
DEVICE = 'mps'
ROOT = r'/Users/mwchang/Downloads/faces'

# Load data 
transformer = transforms.Compose([
    transforms.CenterCrop((128, 128)),
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(), # Scales data into [0,1]
    transforms.Lambda(lambda t: (t * 2) - 1) # Scale between [-1, 1]
])

# Truncate the dataset to a smaller size, so that it's faster to train this example.
data = torchvision.datasets.ImageFolder(root=ROOT, transform=transformer)
data = torch.utils.data.Subset(data, range(SUBSET))
dataloader = DataLoader(data, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)

Initiate Model

#model = UNet(in_channels=IN_CHANNELS, out_channels=IN_CHANNELS, seq_len=T, emb_dim=TIME_EBD_DIM)
model = UNet()
model.to(DEVICE)
diffusion = Diffusion(noise_steps=T, device=DEVICE)

# choose an optimizer
optimizer = AdamW(model.parameters(), lr=0.001)

# choose an loss function
mse = nn.MSELoss()

# Load previous check point if needed
# model, optimizer, epoch = load_checkpoint(model, optimizer, path='model_500.chk')

Print model summary

# Fake data:
x = torch.randn(BATCH_SIZE, IN_CHANNELS, IMG_SIZE, IMG_SIZE).to(DEVICE)
t = torch.randint(0, T, (BATCH_SIZE,)).to(DEVICE)
summary(model, input_data=(x, t), verbose=1)

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
UNet                                     [128, 3, 64, 64]          --
├─Sequential: 1-1                        [128, 256]                --
│    └─PositionalEncoder: 2-1            [128, 256]                --
│    └─Linear: 2-2                       [128, 256]                65,792
│    └─ReLU: 2-3                         [128, 256]                --
├─Block: 1-2                             [128, 64, 64, 64]         --
│    └─Conv2d: 2-4                       [128, 64, 64, 64]         1,792
│    └─ReLU: 2-5                         [128, 64, 64, 64]         --
│    └─BatchNorm2d: 2-6                  [128, 64, 64, 64]         128
│    └─Sequential: 2-7                   [128, 64]                 --
│    │    └─Linear: 3-1                  [128, 64]                 16,448
│    │    └─ReLU: 3-2                    [128, 64]                 --
│    └─Conv2d: 2-8                       [128, 64, 64, 64]         36,928
│    └─ReLU: 2-9                         [128, 64, 64, 64]         --
│    └─BatchNorm2d: 2-10                 [128, 64, 64, 64]         128
├─Down: 1-3                              [128, 128, 32, 32]        --
│    └─Conv2d: 2-11                      [128, 64, 32, 32]         65,600
│    └─Block: 2-12                       [128, 128, 32, 32]        --
│    │    └─Conv2d: 3-3                  [128, 128, 32, 32]        73,856
│    │    └─ReLU: 3-4                    [128, 128, 32, 32]        --
│    │    └─BatchNorm2d: 3-5             [128, 128, 32, 32]        256
│    │    └─Sequential: 3-6              [128, 128]                32,896
│    │    └─Conv2d: 3-7                  [128, 128, 32, 32]        147,584
│    │    └─ReLU: 3-8                    [128, 128, 32, 32]        --
│    │    └─BatchNorm2d: 3-9             [128, 128, 32, 32]        256
├─Down: 1-4                              [128, 256, 16, 16]        --
│    └─Conv2d: 2-13                      [128, 128, 16, 16]        262,272
│    └─Block: 2-14                       [128, 256, 16, 16]        --
│    │    └─Conv2d: 3-10                 [128, 256, 16, 16]        295,168
│    │    └─ReLU: 3-11                   [128, 256, 16, 16]        --
│    │    └─BatchNorm2d: 3-12            [128, 256, 16, 16]        512
│    │    └─Sequential: 3-13             [128, 256]                65,792
│    │    └─Conv2d: 3-14                 [128, 256, 16, 16]        590,080
│    │    └─ReLU: 3-15                   [128, 256, 16, 16]        --
│    │    └─BatchNorm2d: 3-16            [128, 256, 16, 16]        512
├─Down: 1-5                              [128, 512, 8, 8]          --
│    └─Conv2d: 2-15                      [128, 256, 8, 8]          1,048,832
│    └─Block: 2-16                       [128, 512, 8, 8]          --
│    │    └─Conv2d: 3-17                 [128, 512, 8, 8]          1,180,160
│    │    └─ReLU: 3-18                   [128, 512, 8, 8]          --
│    │    └─BatchNorm2d: 3-19            [128, 512, 8, 8]          1,024
│    │    └─Sequential: 3-20             [128, 512]                131,584
│    │    └─Conv2d: 3-21                 [128, 512, 8, 8]          2,359,808
│    │    └─ReLU: 3-22                   [128, 512, 8, 8]          --
│    │    └─BatchNorm2d: 3-23            [128, 512, 8, 8]          1,024
├─Down: 1-6                              [128, 512, 4, 4]          --
│    └─Conv2d: 2-17                      [128, 512, 4, 4]          4,194,816
│    └─Block: 2-18                       [128, 512, 4, 4]          --
│    │    └─Conv2d: 3-24                 [128, 512, 4, 4]          2,359,808
│    │    └─ReLU: 3-25                   [128, 512, 4, 4]          --
│    │    └─BatchNorm2d: 3-26            [128, 512, 4, 4]          1,024
│    │    └─Sequential: 3-27             [128, 512]                131,584
│    │    └─Conv2d: 3-28                 [128, 512, 4, 4]          2,359,808
│    │    └─ReLU: 3-29                   [128, 512, 4, 4]          --
│    │    └─BatchNorm2d: 3-30            [128, 512, 4, 4]          1,024
├─Block: 1-7                             [128, 1024, 4, 4]         --
│    └─Conv2d: 2-19                      [128, 1024, 4, 4]         4,719,616
│    └─ReLU: 2-20                        [128, 1024, 4, 4]         --
│    └─BatchNorm2d: 2-21                 [128, 1024, 4, 4]         2,048
│    └─Sequential: 2-22                  [128, 1024]               --
│    │    └─Linear: 3-31                 [128, 1024]               263,168
│    │    └─ReLU: 3-32                   [128, 1024]               --
│    └─Conv2d: 2-23                      [128, 1024, 4, 4]         9,438,208
│    └─ReLU: 2-24                        [128, 1024, 4, 4]         --
│    └─BatchNorm2d: 2-25                 [128, 1024, 4, 4]         2,048
├─Up: 1-8                                [128, 512, 8, 8]          --
│    └─ConvTranspose2d: 2-26             [128, 512, 8, 8]          2,097,664
│    └─Block: 2-27                       [128, 512, 8, 8]          --
│    │    └─Conv2d: 3-33                 [128, 512, 8, 8]          4,719,104
│    │    └─ReLU: 3-34                   [128, 512, 8, 8]          --
│    │    └─BatchNorm2d: 3-35            [128, 512, 8, 8]          1,024
│    │    └─Sequential: 3-36             [128, 512]                131,584
│    │    └─Conv2d: 3-37                 [128, 512, 8, 8]          2,359,808
│    │    └─ReLU: 3-38                   [128, 512, 8, 8]          --
│    │    └─BatchNorm2d: 3-39            [128, 512, 8, 8]          1,024
├─Up: 1-9                                [128, 256, 16, 16]        --
│    └─ConvTranspose2d: 2-28             [128, 256, 16, 16]        524,544
│    └─Block: 2-29                       [128, 256, 16, 16]        --
│    │    └─Conv2d: 3-40                 [128, 256, 16, 16]        1,179,904
│    │    └─ReLU: 3-41                   [128, 256, 16, 16]        --
│    │    └─BatchNorm2d: 3-42            [128, 256, 16, 16]        512
│    │    └─Sequential: 3-43             [128, 256]                65,792
│    │    └─Conv2d: 3-44                 [128, 256, 16, 16]        590,080
│    │    └─ReLU: 3-45                   [128, 256, 16, 16]        --
│    │    └─BatchNorm2d: 3-46            [128, 256, 16, 16]        512
├─Up: 1-10                               [128, 128, 32, 32]        --
│    └─ConvTranspose2d: 2-30             [128, 128, 32, 32]        131,200
│    └─Block: 2-31                       [128, 128, 32, 32]        --
│    │    └─Conv2d: 3-47                 [128, 128, 32, 32]        295,040
│    │    └─ReLU: 3-48                   [128, 128, 32, 32]        --
│    │    └─BatchNorm2d: 3-49            [128, 128, 32, 32]        256
│    │    └─Sequential: 3-50             [128, 128]                32,896
│    │    └─Conv2d: 3-51                 [128, 128, 32, 32]        147,584
│    │    └─ReLU: 3-52                   [128, 128, 32, 32]        --
│    │    └─BatchNorm2d: 3-53            [128, 128, 32, 32]        256
├─Up: 1-11                               [128, 64, 64, 64]         --
│    └─ConvTranspose2d: 2-32             [128, 64, 64, 64]         32,832
│    └─Block: 2-33                       [128, 64, 64, 64]         --
│    │    └─Conv2d: 3-54                 [128, 64, 64, 64]         73,792
│    │    └─ReLU: 3-55                   [128, 64, 64, 64]         --
│    │    └─BatchNorm2d: 3-56            [128, 64, 64, 64]         128
│    │    └─Sequential: 3-57             [128, 64]                 16,448
│    │    └─Conv2d: 3-58                 [128, 64, 64, 64]         36,928
│    │    └─ReLU: 3-59                   [128, 64, 64, 64]         --
│    │    └─BatchNorm2d: 3-60            [128, 64, 64, 64]         128
├─Block: 1-12                            [128, 3, 64, 64]          --
│    └─Conv2d: 2-34                      [128, 3, 64, 64]          1,731
│    └─ReLU: 2-35                        [128, 3, 64, 64]          --
│    └─BatchNorm2d: 2-36                 [128, 3, 64, 64]          6
│    └─Sequential: 2-37                  [128, 3]                  --
│    │    └─Linear: 3-61                 [128, 3]                  771
│    │    └─ReLU: 3-62                   [128, 3]                  --
│    └─Conv2d: 2-38                      [128, 3, 64, 64]          84
│    └─ReLU: 2-39                        [128, 3, 64, 64]          --
│    └─BatchNorm2d: 2-40                 [128, 3, 64, 64]          6
==========================================================================================
Total params: 42,293,222
Trainable params: 42,293,222
Non-trainable params: 0
Total mult-adds (G): 482.28
==========================================================================================
Input size (MB): 6.29
Forward/backward pass size (MB): 4810.48
Params size (MB): 169.17
Estimated Total Size (MB): 4985.94
==========================================================================================

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
UNet                                     [128, 3, 64, 64]          --
├─Sequential: 1-1                        [128, 256]                --
│    └─PositionalEncoder: 2-1            [128, 256]                --
│    └─Linear: 2-2                       [128, 256]                65,792
│    └─ReLU: 2-3                         [128, 256]                --
├─Block: 1-2                             [128, 64, 64, 64]         --
│    └─Conv2d: 2-4                       [128, 64, 64, 64]         1,792
│    └─ReLU: 2-5                         [128, 64, 64, 64]         --
│    └─BatchNorm2d: 2-6                  [128, 64, 64, 64]         128
│    └─Sequential: 2-7                   [128, 64]                 --
│    │    └─Linear: 3-1                  [128, 64]                 16,448
│    │    └─ReLU: 3-2                    [128, 64]                 --
│    └─Conv2d: 2-8                       [128, 64, 64, 64]         36,928
│    └─ReLU: 2-9                         [128, 64, 64, 64]         --
│    └─BatchNorm2d: 2-10                 [128, 64, 64, 64]         128
├─Down: 1-3                              [128, 128, 32, 32]        --
│    └─Conv2d: 2-11                      [128, 64, 32, 32]         65,600
│    └─Block: 2-12                       [128, 128, 32, 32]        --
│    │    └─Conv2d: 3-3                  [128, 128, 32, 32]        73,856
│    │    └─ReLU: 3-4                    [128, 128, 32, 32]        --
│    │    └─BatchNorm2d: 3-5             [128, 128, 32, 32]        256
│    │    └─Sequential: 3-6              [128, 128]                32,896
│    │    └─Conv2d: 3-7                  [128, 128, 32, 32]        147,584
│    │    └─ReLU: 3-8                    [128, 128, 32, 32]        --
│    │    └─BatchNorm2d: 3-9             [128, 128, 32, 32]        256
├─Down: 1-4                              [128, 256, 16, 16]        --
│    └─Conv2d: 2-13                      [128, 128, 16, 16]        262,272
│    └─Block: 2-14                       [128, 256, 16, 16]        --
│    │    └─Conv2d: 3-10                 [128, 256, 16, 16]        295,168
│    │    └─ReLU: 3-11                   [128, 256, 16, 16]        --
│    │    └─BatchNorm2d: 3-12            [128, 256, 16, 16]        512
│    │    └─Sequential: 3-13             [128, 256]                65,792
│    │    └─Conv2d: 3-14                 [128, 256, 16, 16]        590,080
│    │    └─ReLU: 3-15                   [128, 256, 16, 16]        --
│    │    └─BatchNorm2d: 3-16            [128, 256, 16, 16]        512
├─Down: 1-5                              [128, 512, 8, 8]          --
│    └─Conv2d: 2-15                      [128, 256, 8, 8]          1,048,832
│    └─Block: 2-16                       [128, 512, 8, 8]          --
│    │    └─Conv2d: 3-17                 [128, 512, 8, 8]          1,180,160
│    │    └─ReLU: 3-18                   [128, 512, 8, 8]          --
│    │    └─BatchNorm2d: 3-19            [128, 512, 8, 8]          1,024
│    │    └─Sequential: 3-20             [128, 512]                131,584
│    │    └─Conv2d: 3-21                 [128, 512, 8, 8]          2,359,808
│    │    └─ReLU: 3-22                   [128, 512, 8, 8]          --
│    │    └─BatchNorm2d: 3-23            [128, 512, 8, 8]          1,024
├─Down: 1-6                              [128, 512, 4, 4]          --
│    └─Conv2d: 2-17                      [128, 512, 4, 4]          4,194,816
│    └─Block: 2-18                       [128, 512, 4, 4]          --
│    │    └─Conv2d: 3-24                 [128, 512, 4, 4]          2,359,808
│    │    └─ReLU: 3-25                   [128, 512, 4, 4]          --
│    │    └─BatchNorm2d: 3-26            [128, 512, 4, 4]          1,024
│    │    └─Sequential: 3-27             [128, 512]                131,584
│    │    └─Conv2d: 3-28                 [128, 512, 4, 4]          2,359,808
│    │    └─ReLU: 3-29                   [128, 512, 4, 4]          --
│    │    └─BatchNorm2d: 3-30            [128, 512, 4, 4]          1,024
├─Block: 1-7                             [128, 1024, 4, 4]         --
│    └─Conv2d: 2-19                      [128, 1024, 4, 4]         4,719,616
│    └─ReLU: 2-20                        [128, 1024, 4, 4]         --
│    └─BatchNorm2d: 2-21                 [128, 1024, 4, 4]         2,048
│    └─Sequential: 2-22                  [128, 1024]               --
│    │    └─Linear: 3-31                 [128, 1024]               263,168
│    │    └─ReLU: 3-32                   [128, 1024]               --
│    └─Conv2d: 2-23                      [128, 1024, 4, 4]         9,438,208
│    └─ReLU: 2-24                        [128, 1024, 4, 4]         --
│    └─BatchNorm2d: 2-25                 [128, 1024, 4, 4]         2,048
├─Up: 1-8                                [128, 512, 8, 8]          --
│    └─ConvTranspose2d: 2-26             [128, 512, 8, 8]          2,097,664
│    └─Block: 2-27                       [128, 512, 8, 8]          --
│    │    └─Conv2d: 3-33                 [128, 512, 8, 8]          4,719,104
│    │    └─ReLU: 3-34                   [128, 512, 8, 8]          --
│    │    └─BatchNorm2d: 3-35            [128, 512, 8, 8]          1,024
│    │    └─Sequential: 3-36             [128, 512]                131,584
│    │    └─Conv2d: 3-37                 [128, 512, 8, 8]          2,359,808
│    │    └─ReLU: 3-38                   [128, 512, 8, 8]          --
│    │    └─BatchNorm2d: 3-39            [128, 512, 8, 8]          1,024
├─Up: 1-9                                [128, 256, 16, 16]        --
│    └─ConvTranspose2d: 2-28             [128, 256, 16, 16]        524,544
│    └─Block: 2-29                       [128, 256, 16, 16]        --
│    │    └─Conv2d: 3-40                 [128, 256, 16, 16]        1,179,904
│    │    └─ReLU: 3-41                   [128, 256, 16, 16]        --
│    │    └─BatchNorm2d: 3-42            [128, 256, 16, 16]        512
│    │    └─Sequential: 3-43             [128, 256]                65,792
│    │    └─Conv2d: 3-44                 [128, 256, 16, 16]        590,080
│    │    └─ReLU: 3-45                   [128, 256, 16, 16]        --
│    │    └─BatchNorm2d: 3-46            [128, 256, 16, 16]        512
├─Up: 1-10                               [128, 128, 32, 32]        --
│    └─ConvTranspose2d: 2-30             [128, 128, 32, 32]        131,200
│    └─Block: 2-31                       [128, 128, 32, 32]        --
│    │    └─Conv2d: 3-47                 [128, 128, 32, 32]        295,040
│    │    └─ReLU: 3-48                   [128, 128, 32, 32]        --
│    │    └─BatchNorm2d: 3-49            [128, 128, 32, 32]        256
│    │    └─Sequential: 3-50             [128, 128]                32,896
│    │    └─Conv2d: 3-51                 [128, 128, 32, 32]        147,584
│    │    └─ReLU: 3-52                   [128, 128, 32, 32]        --
│    │    └─BatchNorm2d: 3-53            [128, 128, 32, 32]        256
├─Up: 1-11                               [128, 64, 64, 64]         --
│    └─ConvTranspose2d: 2-32             [128, 64, 64, 64]         32,832
│    └─Block: 2-33                       [128, 64, 64, 64]         --
│    │    └─Conv2d: 3-54                 [128, 64, 64, 64]         73,792
│    │    └─ReLU: 3-55                   [128, 64, 64, 64]         --
│    │    └─BatchNorm2d: 3-56            [128, 64, 64, 64]         128
│    │    └─Sequential: 3-57             [128, 64]                 16,448
│    │    └─Conv2d: 3-58                 [128, 64, 64, 64]         36,928
│    │    └─ReLU: 3-59                   [128, 64, 64, 64]         --
│    │    └─BatchNorm2d: 3-60            [128, 64, 64, 64]         128
├─Block: 1-12                            [128, 3, 64, 64]          --
│    └─Conv2d: 2-34                      [128, 3, 64, 64]          1,731
│    └─ReLU: 2-35                        [128, 3, 64, 64]          --
│    └─BatchNorm2d: 2-36                 [128, 3, 64, 64]          6
│    └─Sequential: 2-37                  [128, 3]                  --
│    │    └─Linear: 3-61                 [128, 3]                  771
│    │    └─ReLU: 3-62                   [128, 3]                  --
│    └─Conv2d: 2-38                      [128, 3, 64, 64]          84
│    └─ReLU: 2-39                        [128, 3, 64, 64]          --
│    └─BatchNorm2d: 2-40                 [128, 3, 64, 64]          6
==========================================================================================
Total params: 42,293,222
Trainable params: 42,293,222
Non-trainable params: 0
Total mult-adds (G): 482.28
==========================================================================================
Input size (MB): 6.29
Forward/backward pass size (MB): 4810.48
Params size (MB): 169.17
Estimated Total Size (MB): 4985.94
==========================================================================================

Test forward diffusion

images, _ = next(iter(dataloader))
images = images.to(DEVICE)  # shape: (N, C, H, W)
img = images[0]  # shape: (C, H, W)
img = img.unsqueeze(0)  # shape: (1, C, H, W)
plt.figure(figsize=(20, 20))

num_images = 10
stepsize = int(T/num_images)

noised_imgs = []
for t in range(0, T, stepsize):
    noised_img, epsilon = diffusion.noise_image_at_time_t(img, t)  # shape: (N, C, H, W)
    noised_imgs.append(noised_img)
    
noised_imgs = torch.cat(noised_imgs, dim=0)  # shape: (N, C, H, W)
plot_tensor_images(noised_imgs, nrow=num_images)

Training

EPOCHS = 500
for epoch in range(201, EPOCHS + 1):
    for step, (images, _) in enumerate(dataloader):
        """This is the implementation of Algorithm 1 in the paper: https://arxiv.org/abs/2006.11239 """
        
        images = images.to(DEVICE)
        t = diffusion.sample_timesteps(n=len(images))
        x_t, noise = diffusion.noise_image_at_time_t(images, t)  # shape: (N, C, H, W); t is (N, 1)
        
        predicted_noise = model(x_t, t)
        loss = mse(predicted_noise, noise)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        print(f"Epoch {epoch} | step {step:03d} Loss: {loss.item()} ")
        if (epoch % 5 == 0 and step == 0) or (epoch == 1 and step == 0):
            save_checkpoint(model, optimizer, epoch, path=f'model.chk')
            sampled_images = diffusion.sample(model, n=5)
            plot_tensor_images(sampled_images, path=f'model_epoch{epoch}.png', nrow=10)

Epoch 500 | step 001 Loss: 0.02067510597407818 
Epoch 500 | step 002 Loss: 0.019939690828323364 
Epoch 500 | step 003 Loss: 0.019561586901545525 
Epoch 500 | step 004 Loss: 0.018764477223157883 
Epoch 500 | step 005 Loss: 0.025068841874599457 
Epoch 500 | step 006 Loss: 0.01874428801238537 
Epoch 500 | step 007 Loss: 0.020428843796253204 
Epoch 500 | step 008 Loss: 0.017568906769156456 
Epoch 500 | step 009 Loss: 0.021536987274885178 
Epoch 500 | step 010 Loss: 0.023131608963012695 
Epoch 500 | step 011 Loss: 0.023605098947882652 
Epoch 500 | step 012 Loss: 0.020491808652877808 
Epoch 500 | step 013 Loss: 0.022295337170362473 
Epoch 500 | step 014 Loss: 0.02512245997786522 
Epoch 500 | step 015 Loss: 0.022736210376024246 
Epoch 500 | step 016 Loss: 0.0201430581510067 
Epoch 500 | step 017 Loss: 0.018344802781939507 
Epoch 500 | step 018 Loss: 0.019374331459403038 
Epoch 500 | step 019 Loss: 0.0240789782255888 
Epoch 500 | step 020 Loss: 0.0158533975481987 
Epoch 500 | step 021 Loss: 0.016398921608924866 
Epoch 500 | step 022 Loss: 0.016088832169771194 
Epoch 500 | step 023 Loss: 0.018646908923983574 
Epoch 500 | step 024 Loss: 0.024842344224452972 
Epoch 500 | step 025 Loss: 0.030117914080619812 
Epoch 500 | step 026 Loss: 0.01823798380792141 
Epoch 500 | step 027 Loss: 0.022577296942472458 
Epoch 500 | step 028 Loss: 0.02231389284133911 
Epoch 500 | step 029 Loss: 0.02090815268456936 
Epoch 500 | step 030 Loss: 0.019084539264440536 
Epoch 500 | step 031 Loss: 0.02585926279425621 
Epoch 500 | step 032 Loss: 0.024560963734984398 
Epoch 500 | step 033 Loss: 0.022622477263212204 
Epoch 500 | step 034 Loss: 0.022576957941055298 
Epoch 500 | step 035 Loss: 0.020859038457274437 
Epoch 500 | step 036 Loss: 0.020285876467823982 
Epoch 500 | step 037 Loss: 0.02503911219537258 
Epoch 500 | step 038 Loss: 0.02849581092596054 
Epoch 500 | step 039 Loss: 0.019079040735960007 
Epoch 500 | step 040 Loss: 0.01662837527692318 
Epoch 500 | step 041 Loss: 0.020989693701267242 
Epoch 500 | step 042 Loss: 0.02171657606959343 
Epoch 500 | step 043 Loss: 0.02710052952170372 
Epoch 500 | step 044 Loss: 0.022889016196131706 
Epoch 500 | step 045 Loss: 0.027236413210630417 
Epoch 500 | step 046 Loss: 0.014890301041305065 
Epoch 500 | step 047 Loss: 0.017803959548473358 
Epoch 500 | step 048 Loss: 0.021787619218230247 
Epoch 500 | step 049 Loss: 0.024534767493605614 
Epoch 500 | step 050 Loss: 0.03111540526151657 
Epoch 500 | step 051 Loss: 0.022318892180919647 
Epoch 500 | step 052 Loss: 0.022493690252304077 
Epoch 500 | step 053 Loss: 0.016964402049779892 
Epoch 500 | step 054 Loss: 0.01811649091541767 
Epoch 500 | step 055 Loss: 0.021130643784999847 
Epoch 500 | step 056 Loss: 0.024370089173316956 
Epoch 500 | step 057 Loss: 0.019411148503422737 
Epoch 500 | step 058 Loss: 0.028876878321170807 
Epoch 500 | step 059 Loss: 0.018544543534517288 
Epoch 500 | step 060 Loss: 0.02236158773303032 
Epoch 500 | step 061 Loss: 0.02667180821299553 
Epoch 500 | step 062 Loss: 0.022771667689085007 
Epoch 500 | step 063 Loss: 0.02462456189095974 
Epoch 500 | step 064 Loss: 0.025312911719083786 
Epoch 500 | step 065 Loss: 0.0235355906188488 
Epoch 500 | step 066 Loss: 0.01852402836084366 
Epoch 500 | step 067 Loss: 0.019029764458537102 
Epoch 500 | step 068 Loss: 0.017660891637206078 
Epoch 500 | step 069 Loss: 0.029713407158851624 
Epoch 500 | step 070 Loss: 0.021472012624144554 
Epoch 500 | step 071 Loss: 0.015576671808958054 
Epoch 500 | step 072 Loss: 0.02374913915991783 
Epoch 500 | step 073 Loss: 0.01850845292210579 
Epoch 500 | step 074 Loss: 0.02009177953004837 
Epoch 500 | step 075 Loss: 0.020707879215478897 
Epoch 500 | step 076 Loss: 0.025493387132883072 
Epoch 500 | step 077 Loss: 0.019601304084062576 
Epoch 500 | step 078 Loss: 0.01753057911992073 
Epoch 500 | step 079 Loss: 0.0339120589196682 
Epoch 500 | step 080 Loss: 0.025775864720344543 
Epoch 500 | step 081 Loss: 0.022769590839743614 
Epoch 500 | step 082 Loss: 0.02093612775206566 
Epoch 500 | step 083 Loss: 0.019365927204489708 
Epoch 500 | step 084 Loss: 0.025675108656287193 
Epoch 500 | step 085 Loss: 0.022980164736509323 
Epoch 500 | step 086 Loss: 0.021466881036758423 
Epoch 500 | step 087 Loss: 0.022218389436602592 
Epoch 500 | step 088 Loss: 0.024551643058657646 
Epoch 500 | step 089 Loss: 0.017855670303106308 
Epoch 500 | step 090 Loss: 0.020540636032819748 
Epoch 500 | step 091 Loss: 0.019024059176445007 
Epoch 500 | step 092 Loss: 0.015634557232260704 
Epoch 500 | step 093 Loss: 0.024687033146619797 
Epoch 500 | step 094 Loss: 0.02068755403161049 
Epoch 500 | step 095 Loss: 0.019301187247037888 
Epoch 500 | step 096 Loss: 0.02319318614900112 
Epoch 500 | step 097 Loss: 0.02349945716559887 
Epoch 500 | step 098 Loss: 0.021571433171629906 
Epoch 500 | step 099 Loss: 0.020617607980966568 
Epoch 500 | step 100 Loss: 0.021936532109975815 
Epoch 500 | step 101 Loss: 0.01804320514202118 
Epoch 500 | step 102 Loss: 0.024707568809390068 
Epoch 500 | step 103 Loss: 0.020859716460108757 
Epoch 500 | step 104 Loss: 0.020980579778552055 
Epoch 500 | step 105 Loss: 0.02153278887271881 
Epoch 500 | step 106 Loss: 0.019736453890800476 
Epoch 500 | step 107 Loss: 0.025488078594207764 
Epoch 500 | step 108 Loss: 0.02347862720489502 
Epoch 500 | step 109 Loss: 0.018598636612296104 
Epoch 500 | step 110 Loss: 0.02101946622133255 
Epoch 500 | step 111 Loss: 0.03516718000173569 
Epoch 500 | step 112 Loss: 0.02316613495349884 
Epoch 500 | step 113 Loss: 0.024032235145568848 
Epoch 500 | step 114 Loss: 0.022251790389418602 
Epoch 500 | step 115 Loss: 0.023179389536380768 
Epoch 500 | step 116 Loss: 0.027668487280607224 
Epoch 500 | step 117 Loss: 0.02824624441564083 
Epoch 500 | step 118 Loss: 0.0231002327054739 
Epoch 500 | step 119 Loss: 0.022393569350242615 
Epoch 500 | step 120 Loss: 0.021673928946256638 
Epoch 500 | step 121 Loss: 0.02351362071931362 
Epoch 500 | step 122 Loss: 0.01953340694308281 
Epoch 500 | step 123 Loss: 0.014722096733748913 
Epoch 500 | step 124 Loss: 0.019460760056972504 
Epoch 500 | step 125 Loss: 0.015004202723503113 
Epoch 500 | step 126 Loss: 0.01823727786540985 
Epoch 500 | step 127 Loss: 0.024626944214105606 
Epoch 500 | step 128 Loss: 0.02018214762210846 
Epoch 500 | step 129 Loss: 0.023589100688695908 
Epoch 500 | step 130 Loss: 0.025462036952376366 
Epoch 500 | step 131 Loss: 0.021607739850878716 
Epoch 500 | step 132 Loss: 0.02213825285434723 
Epoch 500 | step 133 Loss: 0.026319976896047592 
Epoch 500 | step 134 Loss: 0.02284703217446804 
Epoch 500 | step 135 Loss: 0.023279158398509026 
Epoch 500 | step 136 Loss: 0.016467846930027008 
Epoch 500 | step 137 Loss: 0.023183640092611313 
Epoch 500 | step 138 Loss: 0.020428404211997986 
Epoch 500 | step 139 Loss: 0.027814023196697235 
Epoch 500 | step 140 Loss: 0.016164477914571762 
Epoch 500 | step 141 Loss: 0.01807009056210518 
Epoch 500 | step 142 Loss: 0.02308664284646511 
Epoch 500 | step 143 Loss: 0.020361728966236115 
Epoch 500 | step 144 Loss: 0.020902510732412338 
Epoch 500 | step 145 Loss: 0.020814433693885803 
Epoch 500 | step 146 Loss: 0.017225421965122223 
Epoch 500 | step 147 Loss: 0.023127347230911255 
Epoch 500 | step 148 Loss: 0.02370229922235012 
Epoch 500 | step 149 Loss: 0.01986108534038067 
Epoch 500 | step 150 Loss: 0.02225065976381302 
Epoch 500 | step 151 Loss: 0.018742050975561142 
Epoch 500 | step 152 Loss: 0.022652268409729004 
Epoch 500 | step 153 Loss: 0.017585765570402145 
Epoch 500 | step 154 Loss: 0.023062175139784813 
Epoch 500 | step 155 Loss: 0.018724475055933 

Snapshots of training

Plot sample images

sampled_images = diffusion.sample(model, n=100)
plt.figure(figsize=(16,16))
plot_tensor_images(sampled_images, nrow=10)

999it [02:49,  5.90it/s]

save_checkpoint(model, optimizer, epoch, path=f'model_500.chk')